// Copyright 2022 jeff.li. and/or its affiliates.

#pragma once

#include <memory>

#include <tbir/pipeline/op_kernel.h>
#include <tbir/runtime/threadpool/i_thread_pool.h>

namespace tbir::runtime {

    class ThreadPoolOp : public OpKernel {
    public:
        // Before fork, the parent process should call this function once
        void AtForkBefore();

        // After fork, the child/parent process should call this function once
        void AtForkAfterInParentOrChild();

        void Init() override;

        const std::shared_ptr<internal::IThreadPool> &GetPool() const {
            return pool_;
        }

        RTValue Process(PyArgs inputs) const override;

    private:
        bool lock_free_ = false;
        int32_t thread_nums_ = false;
        Unicode thread_name_;
        std::shared_ptr<internal::IThreadPool> pool_ = nullptr;
    };
}  // namespace tbir::runtime